"""
SELECT TABLE_NAME AS `Table Name`, TABLE_COMMENT AS `Table Comment`
FROM information_schema.TABLES
WHERE TABLE_SCHEMA = 'student';
"""

from langchain import SQLDatabase

# 配置数据库连接信息

db_user = "sto"
db_password = "Sto@2025"
db_host = "127.0.0.1"
db_name = "sto"
"""
db_user = "dev"
db_password = "123456"
db_host = "192.168.28.174"
db_name = "cloud_desk"
"""
# 构建数据库连接 URI
db_uri = f"mysql+mysqlconnector://{db_user}:{db_password}@{db_host}/{db_name}"

# 创建数据库连接对象
db = SQLDatabase.from_uri(db_uri)
print(db)

import requests
from langchain.llms.base import LLM
from typing import Optional, List, Mapping, Any

from openai import OpenAI

from typing import Iterator
from langchain_core.outputs import GenerationChunk
from langchain_core.callbacks.manager import CallbackManagerForLLMRun

class DeepSeekLLM(LLM):
    """DeepSeek 语言模型"""

    @property
    def _llm_type(self) -> str:
        return "deepseek"

    def _call(
        self,
        prompt: str,
        stop: Optional[List[str]] = None,
        run_manager: Optional[CallbackManagerForLLMRun] = None,
        **kwargs: Any,
    ) -> str:
        client = OpenAI(
            base_url="https://api.siliconflow.cn/v1",
            api_key="sk-gskfqqyvcayxexdtehxwmhjtlkfqqhygoffgunzvpqckdkov",
        )

        chat_completion = client.chat.completions.create(
            messages=[
                {
                    "role": "user",
                    "content": prompt,
                }
            ],
            # model='deepseek-ai/DeepSeek-V3',
            # model='deepseek-ai/DeepSeek-R1',
            # model="deepseek-ai/DeepSeek-R1-Distill-Llama-8B",
            # model="deepseek-ai/DeepSeek-R1-Distill-Qwen-7B",
            model="deepseek-ai/DeepSeek-R1-Distill-Qwen-7B",
            temperature=0.2,
        )
        if self.verbose:
            print(chat_completion.choices[0].message.content + "\n")
            print(f"DeepSeek LLM: Prompt Tokens: {chat_completion.usage.prompt_tokens} | Response Tokens: {chat_completion.usage.completion_tokens }\n")
        return chat_completion.choices[0].message.content

    def _stream(
        self,
        prompt: str,
        stop: Optional[List[str]] = None,
        run_manager: Optional[CallbackManagerForLLMRun] = None,
        **kwargs: Any,
    ) -> Iterator[GenerationChunk]:
        for char in prompt:
            chunk = GenerationChunk(text=char)
            if run_manager:
                run_manager.on_llm_new_token(chunk.text, chunk=chunk)

            yield chunk

    @property
    def _identifying_params(self) -> Mapping[str, Any]:
        return {}

# 创建 DeepSeek 语言模型实例
llm = DeepSeekLLM(verbose=True)

from langchain_community.utilities import SQLDatabase
from langchain.chains.llm import LLMChain
from langchain_experimental.sql import SQLDatabaseSequentialChain, SQLDatabaseChain
from langchain.llms.base import LLM
from langchain.prompts.prompt import PromptTemplate

# 定义自定义模板
TEMPLATE = """
Given an input question, first create a syntactically correct MySQL query to run, then look at the results of the query and return the answer to the input question.
Unless the user specifies in the question a specific number of examples to obtain.
You can order the results to make them more informative.
Never query for all columns from a table. Instead, specify only the columns you need.
Pay attention to use only the column names that you can see in the schema. Be careful to not query for columns that do not exist.
Pay attention to which column is in which table.
If you get an error while executing a query, rewrite the query and try again.

基于已知信息进行分析，简洁和专业的来回答用户的问题。如果无法从中得到答案，请说 "根据已知信息无法回答该问题"，不允许猜测已知信息意外的表，答案请使用中文。

已知信息：
ab1_1	
ab2_1	
account_detail	用户账户
account_exchange	股票历史交易
acts	自选股
auto_select	机器选股表
c_1	
code_news	股市新闻表
codes	股市代码表
codes_metrics	股票统计表
company	公司信息
company_finance_base	基本财务表
company_finance_earn	盈利能力财务表
company_finance_hold	运营能力能力财务表
company_finance_money	现金流量
company_finance_other	财务其它指标
company_finance_return	偿债及资本结构
company_finance_up	成长能力财务表
company_person	股东表
day_kline	股票全市场日K线数据分数据0
day_kline_0	股票全市场日K线数据分数据0
day_kline_1	股票全市场日K线数据分数据0
day_kline_2	股票全市场日K线数据分数据0
day_kline_3	股票全市场日K线数据分数据0
day_kline_4	股票全市场日K线数据分数据0
day_kline_5	股票全市场日K线数据分数据0
day_kline_6	股票全市场日K线数据分数据0
day_kline_7	股票全市场日K线数据分数据0
day_kline_8	股票全市场日K线数据分数据0
day_kline_9	股票全市场日K线数据分数据0
day_kline_all	股票全市场日K线数据分数据0
exam_paper	试卷表
exam_paper_question	试卷-题目关联表
exam_record	考试记录表
index_menu	
jlr_day	每日净流入表
ke4	
menu	
messages	预警信息
money_per_day	每日净收入表
news	资讯信息
question_group	
questions	试卷
questions_jiashizheng_ab1_1	试卷
questions_jiashizheng_ab2_1	试卷
questions_jiashizheng_c1	试卷
questions_jiashizheng_ke4	试卷
role	角色表
stock_alert	预警股票
stock_holds	股票持仓
stock_holds_detail	股票持仓
stock_select	自选股
sys_permission	
talk_day	盘面评价表
talk_user_day	个人持股评价表
task	任务执行情况表
user	
user_answer	用户答题记录表
users  
zhangting	涨停统计表
zhangting_bk	涨停统计表
zs_info_day	每日指数信息表

Must Use the following format for the output with out any extra information or markdown formatting:

Question: {input}
Answer: Final answer here"""

from langchain_core.output_parsers.list import CommaSeparatedListOutputParser  # noqa: E402

# 创建 PromptTemplate 对象
PROMPT = PromptTemplate(
    input_variables=["input"],
    template=TEMPLATE,
    output_parser=CommaSeparatedListOutputParser(),
)

# 创建 SQLDatabaseChain
db_chain = SQLDatabaseChain.from_llm(llm, db, verbose=True)

# 示例问题
question = "股票代码SH002024的股东名单及占股比例?"

# 运行查询
result = db_chain(question)

print("提取结果字符串:")
print(result)

